import torch
from torch import Tensor, nn
from torch.nn import Module
from typing import Union, List, Tuple, Iterable, Optional, Dict, Any
from datetime import time, timedelta, date
import numpy as np
from numpy import ndarray
import yaml
import os
from functools import lru_cache
from pathlib import Path

_cur_dir = Path.cwd()
MAX_PATH_LENGTH = 1024

def pbc_box_reshape(pbc_box: Tensor, ndim: int) -> Tensor:
    """Reshape the pbc_box as the same ndim with JIT optimization"""
    if ndim <= 2:
        return pbc_box
    shape = pbc_box.shape[:1] + (1,) * (ndim - 2) + pbc_box.shape[-1:]
    return torch.reshape(pbc_box, shape)

def vector_in_pbc(vector: Tensor, pbc_box: Tensor, offset: float = -0.5) -> Tensor:
    """Make the value of vector in a single PBC box with JIT optimization"""
    pbc_box = pbc_box_reshape(pbc_box, vector.ndim)
    box_nograd = pbc_box.detach()
    inv_box = torch.reciprocal(box_nograd)
    vector -= box_nograd * torch.floor(vector * inv_box - offset)
    return vector * inv_box * pbc_box

def calc_vector_pbc(initial: Tensor, terminal: Tensor, pbc_box: Tensor) -> Tensor:
    """Calculate vector with PBC with JIT optimization"""
    return vector_in_pbc(terminal-initial, pbc_box)

def calc_vector(initial: Tensor, terminal: Tensor, pbc_box: Optional[Tensor] = None) -> Tensor:
    """Calculate vector with optional PBC with JIT optimization"""
    if initial is None or terminal is None:
        raise ValueError("initial and terminal tensors cannot be None")
    vector = terminal - initial
    if pbc_box is None or pbc_box.numel() == 0:
        return vector
    return vector_in_pbc(vector, pbc_box)

class GetVector(Module):
    """Vector calculation module with improved performance"""

    def __init__(self, use_pbc: Optional[bool] = None):
        super().__init__()
        self._use_pbc = use_pbc
        self._set_calc_method()

    def _set_calc_method(self):
        """Set calculation method based on PBC flag"""
        if self._use_pbc is None:
            self.calc_vector = self.calc_vector_default
        else:
            self.calc_vector = self.calc_vector_pbc if self._use_pbc else self.calc_vector_nopbc

    @property
    def use_pbc(self) -> Optional[bool]:
        """Get PBC usage flag"""
        return self._use_pbc

    @use_pbc.setter
    def use_pbc(self, use_pbc_: bool):
        """Set PBC usage flag"""
        self.set_pbc(use_pbc_)

    def set_pbc(self, use_pbc: bool):
        """Set PBC usage with validation"""
        if use_pbc is not None and not isinstance(use_pbc, bool):
            raise TypeError(f'use_pbc must be bool or None but got: {type(use_pbc)}')
        self._use_pbc = use_pbc
        self._set_calc_method()
        return self

    def calc_vector_default(self, initial: Tensor, terminal: Tensor, 
                          pbc_box: Optional[Tensor] = None) -> Tensor:
        """Default vector calculation with JIT optimization"""
        if initial is None or terminal is None:
            raise ValueError("initial and terminal tensors cannot be None")
        vector = terminal - initial
        if pbc_box is None or pbc_box.numel() == 0:
            return vector
        return vector_in_pbc(vector, pbc_box)

    def calc_vector_pbc(self, initial: Tensor, terminal: Tensor,
                       pbc_box: Optional[Tensor] = None) -> Tensor:
        """PBC vector calculation with JIT optimization"""
        if initial is None or terminal is None:
            raise ValueError("initial and terminal tensors cannot be None")
        if pbc_box is None or pbc_box.numel() == 0:
            return terminal - initial
        return vector_in_pbc(terminal - initial, pbc_box)

    def calc_vector_nopbc(self, initial: Tensor, terminal: Tensor,
                         pbc_box: Optional[Tensor] = None) -> Tensor:
        """Non-PBC vector calculation with JIT optimization"""
        if initial is None or terminal is None:
            raise ValueError("initial and terminal tensors cannot be None")
        return terminal - initial

    def forward(self, initial: Tensor, terminal: Tensor, 
               pbc_box: Optional[Tensor] = None) -> Tensor:
        """Forward pass with improved type hints"""
        return self.calc_vector(initial, terminal, pbc_box)

def get_arguments(locals_: dict, kwargs: Optional[dict] = None) -> Dict[str, Any]:
    """Get arguments with caching"""
    if '__class__' in locals_:
        locals_.pop('__class__')

    arguments: Dict[str, Any] = {}
    if 'self' in locals_:
        cls = locals_.pop('self')
        arguments['cls_name'] = cls.__class__.__name__

    def _convert(value: Any) -> Any:
        """Convert value with type checking"""
        if value is None or isinstance(value, (int, float, bool, str,
                                           time, timedelta, date)):
            return value
        if isinstance(value, ndarray):
            return value.tolist()
        if isinstance(value, (Tensor, nn.Parameter)):
            return value.cpu().numpy().tolist()
        if isinstance(value, (list, tuple)):
            return [_convert(v) for v in value]
        if isinstance(value, dict):
            if 'cls_name' in value:
                return value
            return {k: _convert(v) for k, v in value.items()}

        cls_name = value.__class__.__name__
        if hasattr(value, '_kwargs'):
            value = value.__dict__['_kwargs']
        elif hasattr(value, 'init_args'):
            value = value.__dict__['init_args']
        else:
            value = value.__class__.__name__

        if isinstance(value, dict) and 'cls_name' not in value:
            value = {'cls_name': cls_name, **_convert(value)}

        return value

    kwargs_ = locals_.pop('kwargs', {}) if 'kwargs' in locals_ else {}
    kwargs = kwargs or kwargs_
    if 'cls_name' in kwargs:
        kwargs.pop('cls_name')

    arguments.update(_convert(locals_))
    arguments.update(_convert(kwargs))

    return arguments

class CybertorchDevice:
    def __init__(self, device) -> None:
        self.device = device

    def set_device(self, device):
        self.device = device
    def __call__(self):
        return self.device

GLOBAL_DEVICE = CybertorchDevice('cpu')

def set_cybertorch_global_device(device):
    global GLOBAL_DEVICE
    GLOBAL_DEVICE.set_device(device)

def get_tensor(value: Union[float, int, Tensor, nn.Parameter, ndarray, List[float], Tuple[float]],
               dtype: Optional[torch.dtype] = None,
               device: Union[str, torch.device] = 'cpu') -> Optional[Tensor]:
    """Get tensor with improved type handling"""
    if value is None:
        return None

    # Convert device to torch.device if needed
    if isinstance(device, str):
        device = torch.device(device)

    if isinstance(value, ndarray):
        return torch.as_tensor(value, dtype=dtype, device=device)
    if isinstance(value, nn.Parameter):
        return value.clone().detach().to(dtype=dtype, device=device)
    if not isinstance(value, Tensor):
        return torch.tensor([value], dtype=dtype, device=device)[0]
    return value.to(dtype=dtype, device=device) if dtype else value

def get_integer(value: Union[int, Tensor, nn.Parameter, ndarray]) -> Optional[int]:
    """Get integer with improved type handling"""
    if value is None:
        return None
    if isinstance(value, Tensor):
        value = value.cpu().numpy()
    return int(value)

INITIALIZERS = {name for name in dir(torch.nn.init) if not name.startswith('_')}

def get_initializer(cls_name: Union[str, Tensor, Dict[str, Any]], 
                   in_features: int,
                   out_features: Optional[int] = None,
                   dtype: torch.dtype = torch.float32,
                   device: str = 'cpu',
                   **kwargs) -> Tensor:
    """Get initializer with improved error handling"""
    if isinstance(cls_name, str):
        init_name = cls_name if cls_name.endswith('_') else f"{cls_name}_"
        if init_name.lower() not in INITIALIZERS:
            raise ValueError(f"Initializer '{cls_name}' not found")
        init_fn = getattr(torch.nn.init, init_name)
        shape = (in_features, out_features) if out_features else (in_features,)
        return init_fn(torch.empty(*shape, dtype=dtype, device=device))

    if isinstance(cls_name, (Tensor, nn.Parameter, ndarray)):
        return get_tensor(cls_name, dtype=dtype)

    if isinstance(cls_name, dict):
        return get_initializer(**cls_name)

    raise TypeError(f'cls_name must be str, dict or Tensor but got: {type(cls_name)}')

def _make_dir(path: Union[str, Path], arg_name: str = 'path') -> Path:
    """Make directory with improved path handling"""
    if not isinstance(path, (str, Path)):
        raise TypeError(f"{arg_name} must be string or Path")
    if isinstance(path, str) and not path.strip():
        raise ValueError(f"{arg_name} must be non-blank")

    path = Path(path).resolve()
    if len(str(path)) > MAX_PATH_LENGTH:
        raise ValueError(f"{arg_name} exceeds maximum length of {MAX_PATH_LENGTH}")

    if path.exists():
        if not path.is_dir():
            raise NotADirectoryError(f"{path} is a file, not a directory")
    else:
        try:
            permissions = os.R_OK | os.W_OK | os.X_OK
            os.umask(permissions << 3 | permissions)
            path.mkdir(mode=permissions << 6, parents=True, exist_ok=True)
        except PermissionError as e:
            raise PermissionError("No write permission") from e

    return path

def read_yaml(filename: Union[str, Path]) -> Dict[str, Any]:
    """Read YAML with improved error handling"""
    try:
        with open(filename, 'r', encoding="utf-8") as f:
            return yaml.safe_load(f)
    except yaml.YAMLError as e:
        raise ValueError(f"Invalid YAML file: {filename}") from e
    except OSError as e:
        raise OSError(f"Failed to read file: {filename}") from e

def write_yaml(data: Dict[str, Any], filename: Union[str, Path], 
              directory: Optional[Union[str, Path]] = None):
    """Write YAML with improved error handling"""
    directory = _make_dir(directory or _cur_dir)
    filepath = Path(directory) / filename
    
    try:
        with open(filepath, 'w', encoding="utf-8") as f:
            yaml.dump(data, f, sort_keys=False)
    except yaml.YAMLError as e:
        raise ValueError(f"Failed to serialize data to YAML") from e
    except OSError as e:
        raise OSError(f"Failed to write file: {filepath}") from e
    
# 设置随机种子
def set_seed(seed1,seed2,seed3):
    torch.manual_seed(seed1)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed2)
    np.random.seed(seed3)